import matplotlib.pyplot as plt
import numpy as np
import json


class SudokuCell:
    def __init__(self, row, col):
        self.row = row
        self.col = col
        self.value = None  # Player entered value
        self.cue = None  # Puzzle clue number
        self.block_color = None
        self.candidates = {}

    def set_value(self, value):
        self.value = value
        self.candidates = set()  # Clear candidates after setting value

    def set_cue(self, cue):
        self.cue = cue
        self.candidates = set()  # Clear candidates after setting clue

    def set_candidate(self, num, color=None):
        if not self.value and not self.cue and 1 <= num <= 9:
            self.candidates[num] = color

    def clear_candidate(self, num):
        if num in self.candidates:
            del self.candidates[num]

    def set_block_color(self, color):
        self.block_color = color


class SudokuBoard:
    def __init__(self):
        self.cells = [
            [SudokuCell(row, col) for col in range(1, 10)] for row in range(1, 10)
        ]

    def set_value(self, row, col, value):
        self.cells[row - 1][col - 1].set_value(value)

    def set_cue(self, row, col, cue):
        self.cells[row - 1][col - 1].set_cue(cue)

    def set_candidate(self, row, col, candidate, color=None):
        self.cells[row - 1][col - 1].set_candidate(candidate, color)

    def set_block_color(self, row, col, color):
        self.cells[row - 1][col - 1].set_block_color(color)

    def draw(self, filename="sudoku_grid.png", dpi=300):
        # 计算需要的figsize来保证1080像素输出
        desired_pixels = 1080
        figsize = desired_pixels / dpi  # 转换为英寸

        fig, ax = plt.subplots(figsize=(figsize, figsize))
        ax.set_xticks([])
        ax.set_yticks([])

        # Add row numbers (left side)
        for row in range(1, 10):
            ax.text(
                -0.3,
                9.5 - row,
                str(row),
                fontsize=12,
                ha="center",
                va="center",
                color="lightgray",
            )

        # Add row numbers (right side)
        for row in range(1, 10):
            ax.text(
                9.3,
                9.5 - row,
                str(row),
                fontsize=12,
                ha="center",
                va="center",
                color="lightgray",
            )

        # Add column numbers (top)
        for col in range(1, 10):
            ax.text(
                col - 0.5,
                9.3,
                str(col),
                fontsize=12,
                ha="center",
                va="center",
                color="lightgray",
            )

        # Add column numbers (bottom)
        for col in range(1, 10):
            ax.text(
                col - 0.5,
                -0.3,
                str(col),
                fontsize=12,
                ha="center",
                va="center",
                color="lightgray",
            )

        # Draw cell values and candidates
        for row in self.cells:
            for cell in row:
                if cell.block_color:
                    rect = plt.Rectangle(
                        (cell.col - 1, 9 - cell.row),
                        1,
                        1,
                        facecolor=cell.block_color,
                        alpha=1.0,
                        edgecolor="none",
                    )
                    ax.add_patch(rect)
                if cell.cue:  # Clue number
                    ax.text(
                        cell.col - 0.5,
                        9.5 - cell.row,
                        str(cell.cue),
                        fontsize=32,
                        ha="center",
                        va="center",
                        weight="normal",
                        color="blue",
                    )
                elif cell.value:  # Entered value
                    ax.text(
                        cell.col - 0.5,
                        9.5 - cell.row,
                        str(cell.value),
                        fontsize=32,
                        ha="center",
                        va="center",
                        weight="normal",
                        color="black",
                    )
                elif cell.candidates:  # Candidate numbers
                    for num, color in cell.candidates.items():
                        # Calculate sub-cell position
                        sub_col = (num - 1) % 3
                        sub_row = 2 - (num - 1) // 3
                        x = cell.col - 1 + 0.33 * sub_col + 0.165
                        y = 9 - cell.row + 0.33 * sub_row + 0.165

                        if color:
                            circle = plt.Circle((x, y), 0.12, color=color, alpha=1.0)
                            ax.add_patch(circle)

                        # Draw candidate number
                        ax.text(
                            x,
                            y,
                            str(num),
                            fontsize=10,
                            ha="center",
                            va="center",
                            color="black",  # Number color fixed as black
                        )

        # Draw grid lines
        for i in range(10):
            if i == 0 or i == 9:
                linewidth = 3
                linecolor = "black"
            elif i % 3 == 0:
                linewidth = 3
                linecolor = "black"
            else:
                linewidth = 1
                linecolor = "black"
            ax.axhline(i, color=linecolor, linewidth=linewidth)
            ax.axvline(i, color=linecolor, linewidth=linewidth)

        ax.set_xlim(-0.02, 9.02)
        ax.set_ylim(-0.02, 9.02)
        plt.savefig(filename, dpi=dpi, bbox_inches="tight", pad_inches=0.1)
        plt.close()


if __name__ == "__main__":
    board = SudokuBoard()
    board.set_cue(1, 1, 5)  # Set R1C1 clue number to 5
    board.set_value(2, 1, 3)  # Player enters 3 in R2C1 (won't override clue)
    board.set_candidate(2, 2, 1, None)  # Set R2C2 candidates to 1,5,9
    board.set_candidate(2, 2, 5, None)
    board.set_candidate(2, 2, 9, "red")
    board.set_candidate(3, 5, 7, "green")
    board.set_block_color(2, 2, "lightblue")  # Set R2C2 background color
    board.draw()
